import os
import os.path as osp
import numpy as np
import random
import collections
import torch
import torchvision
import cv2
from torch.utils import data
import matplotlib.pyplot as plt
from options.train_options import TrainOptions


class VOCDataSet(data.Dataset):
    def __init__(self, opt):
        self.root = opt.dataroot
        self.list_path = opt.train_list_path
        self.crop_h, self.crop_w = opt.crop_size_h, opt.crop_size_w
        self.scale = opt.scale
        self.ignore_label = opt.ignore_label
        self.mean = (128, 128, 128)
        self.is_mirror = opt.mirror
        self.max_iters = opt.max_iters

        # self.mean_bgr = np.array([104.00698793, 116.66876762, 122.67891434])
        self.img_ids = [i_id.strip() for i_id in open(self.list_path)]
        if not self.max_iters == None:
            self.img_ids = self.img_ids * int(np.ceil(float(self.max_iters) / len(self.img_ids)))
        self.files = []
        # for split in ["train", "trainval", "val"]:
        for name in self.img_ids:
            img_file = osp.join(self.root, "JPEGImages/%s.jpg" % name)
            label_file = osp.join(self.root, "SegmentationClassAug/%s.png" % name)
            self.files.append({
                "img": img_file,
                "label": label_file,
                "name": name
            })

    def __len__(self):
        return len(self.files)

    def generate_scale_label(self, image, label):
        #f_scale = 0.5 + random.randint(0, 11) / 10.0
        image = cv2.resize(image, (528,528), interpolation=cv2.INTER_LINEAR)
        label = cv2.resize(label, (528,528), interpolation=cv2.INTER_NEAREST)
        return image, label

    def __getitem__(self, index):
        datafiles = self.files[index]
        image = cv2.imread(datafiles["img"], cv2.IMREAD_COLOR)
        label = cv2.imread(datafiles["label"], cv2.IMREAD_GRAYSCALE)
        size = image.shape
        name = datafiles["name"]
        if self.scale:
            image, label = self.generate_scale_label(image, label)
        image = np.asarray(image, np.float32)
        image -= self.mean
        # img_h, img_w = label.shape
        # pad_h = max(self.crop_h - img_h, 0)
        # pad_w = max(self.crop_w - img_w, 0)
        # if pad_h > 0 or pad_w > 0:
        #     img_pad = cv2.copyMakeBorder(image, 0, pad_h, 0,
        #                                  pad_w, cv2.BORDER_CONSTANT,
        #                                  value=(0.0, 0.0, 0.0))
        #     label_pad = cv2.copyMakeBorder(label, 0, pad_h, 0,
        #                                    pad_w, cv2.BORDER_CONSTANT,
        #                                    value=(self.ignore_label,))
        # else:
        #     img_pad, label_pad = image, label
        #
        # img_h, img_w = label_pad.shape
        # h_off = random.randint(0, img_h - self.crop_h)
        # w_off = random.randint(0, img_w - self.crop_w)
        # # roi = cv2.Rect(w_off, h_off, self.crop_w, self.crop_h);
        # image = np.asarray(img_pad[h_off: h_off + self.crop_h, w_off: w_off + self.crop_w], np.float32)
        # label = np.asarray(label_pad[h_off: h_off + self.crop_h, w_off: w_off + self.crop_w], np.float32)
        # # image = image[:, :, ::-1]  # change to BGR
        image = image.transpose((2, 0, 1))
        if self.is_mirror:
            flip = np.random.choice(2) * 2 - 1
            image = image[:, :, ::flip]
            label = label[:, ::flip]

        return image.copy(), label.copy() #, np.array(size), name
    def name(self):
        name_str = 'VOC12_set'
        return name_str


class VOCDataValSet(data.Dataset):
    def __init__(self, opt):
        self.root = opt.dataroot
        self.list_path = opt.val_list_path
        self.crop_h, self.crop_w = opt.crop_size_h, opt.crop_size_w
        self.scale = opt.scale
        self.ignore_label = opt.ignore_label
        self.mean = (128, 128, 128)
        self.is_mirror = opt.mirror
        self.max_iters = opt.max_iters

        # self.mean_bgr = np.array([104.00698793, 116.66876762, 122.67891434])
        self.img_ids = [i_id.strip() for i_id in open(self.list_path)]
        if not self.max_iters == None:
            self.img_ids = self.img_ids * int(np.ceil(float(self.max_iters) / len(self.img_ids)))
        self.files = []
        # for split in ["train", "trainval", "val"]:
        for name in self.img_ids:
            img_file = osp.join(self.root, "JPEGImages/%s.jpg" % name)
            label_file = osp.join(self.root, "SegmentationClassAug/%s.png" % name)
            self.files.append({
                "img": img_file,
                "label": label_file,
                "name": name
            })

    def __len__(self):
        return len(self.files)

    def generate_scale_label(self, image, label):
        #f_scale = 0.5 + random.randint(0, 11) / 10.0
        image = cv2.resize(image, (528,528), interpolation=cv2.INTER_LINEAR)
        label = cv2.resize(label, (528,528), interpolation=cv2.INTER_NEAREST)
        return image, label

    def __getitem__(self, index):
        datafiles = self.files[index]
        image = cv2.imread(datafiles["img"], cv2.IMREAD_COLOR)
        label = cv2.imread(datafiles["label"], cv2.IMREAD_GRAYSCALE)
        size = image.shape
        name = datafiles["name"]
        if self.scale:
            image, label = self.generate_scale_label(image, label)
        image = np.asarray(image, np.float32)
        image -= self.mean
        # img_h, img_w = label.shape
        # pad_h = max(self.crop_h - img_h, 0)
        # pad_w = max(self.crop_w - img_w, 0)
        # if pad_h > 0 or pad_w > 0:
        #     img_pad = cv2.copyMakeBorder(image, 0, pad_h, 0,
        #                                  pad_w, cv2.BORDER_CONSTANT,
        #                                  value=(0.0, 0.0, 0.0))
        #     label_pad = cv2.copyMakeBorder(label, 0, pad_h, 0,
        #                                    pad_w, cv2.BORDER_CONSTANT,
        #                                    value=(self.ignore_label,))
        # else:
        #     img_pad, label_pad = image, label
        #
        # img_h, img_w = label_pad.shape
        # h_off = random.randint(0, img_h - self.crop_h)
        # w_off = random.randint(0, img_w - self.crop_w)
        # # roi = cv2.Rect(w_off, h_off, self.crop_w, self.crop_h);
        # image = np.asarray(img_pad[h_off: h_off + self.crop_h, w_off: w_off + self.crop_w], np.float32)
        # label = np.asarray(label_pad[h_off: h_off + self.crop_h, w_off: w_off + self.crop_w], np.float32)
        # # image = image[:, :, ::-1]  # change to BGR
        image = image.transpose((2, 0, 1))
        if self.is_mirror:
            flip = np.random.choice(2) * 2 - 1
            image = image[:, :, ::flip]
            label = label[:, ::flip]

        return image.copy(), label.copy() #, np.array(size), name
    def name(self):
        name_str = 'VOC12_val_set'
        return name_str